#ifndef REDBLACKTREE_H
#define REDBLACKTREE_H

#include <iostream>

enum Color { RED, BLACK };

template<typename T>
struct Node {
    T data;
    Color color;
    Node* left;
    Node* right;
    Node* parent;
    
    Node(T value) : data(value), color(RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

template<typename T>
class RedBlackTree {
private:
    Node<T>* root;
    Node<T>* nil;
    
    void leftRotate(Node<T>* x);
    void rightRotate(Node<T>* x);
    void insertFixup(Node<T>* k);
    void deleteFixup(Node<T>* x);
    void transplant(Node<T>* u, Node<T>* v);
    Node<T>* minimum(Node<T>* node);
    void inorderHelper(Node<T>* node) const;
    Node<T>* searchHelper(Node<T>* node, T key) const;
    
public:
    RedBlackTree();
    ~RedBlackTree();
    void insert(T key);
    void remove(T key);
    bool search(T key) const;
    void print() const;
    void clear();
    
private:
    void clearHelper(Node<T>* node);
};

template<typename T>
RedBlackTree<T>::RedBlackTree() {
    nil = new Node<T>(T());
    nil->color = BLACK;
    root = nil;
}

template<typename T>
RedBlackTree<T>::~RedBlackTree() {
    clear();
    delete nil;
}

template<typename T>
void RedBlackTree<T>::leftRotate(Node<T>* x) {
    Node<T>* y = x->right;
    x->right = y->left;
    if (y->left != nil) {
        y->left->parent = x;
    }
    y->parent = x->parent;
    if (x->parent == nullptr) {
        root = y;
    } else if (x == x->parent->left) {
        x->parent->left = y;
    } else {
        x->parent->right = y;
    }
    y->left = x;
    x->parent = y;
}

template<typename T>
void RedBlackTree<T>::rightRotate(Node<T>* x) {
    Node<T>* y = x->left;
    x->left = y->right;
    if (y->right != nil) {
        y->right->parent = x;
    }
    y->parent = x->parent;
    if (x->parent == nullptr) {
        root = y;
    } else if (x == x->parent->right) {
        x->parent->right = y;
    } else {
        x->parent->left = y;
    }
    y->right = x;
    x->parent = y;
}

template<typename T>
void RedBlackTree<T>::insert(T key) {
    Node<T>* z = new Node<T>(key);
    Node<T>* y = nullptr;
    Node<T>* x = root;
    
    while (x != nil) {
        y = x;
        if (z->data < x->data) {
            x = x->left;
        } else {
            x = x->right;
        }
    }
    
    z->parent = y;
    if (y == nullptr) {
        root = z;
    } else if (z->data < y->data) {
        y->left = z;
    } else {
        y->right = z;
    }
    
    z->left = nil;
    z->right = nil;
    z->color = RED;
    
    insertFixup(z);
}

template<typename T>
void RedBlackTree<T>::insertFixup(Node<T>* k) {
    Node<T>* u;
    while (k->parent != nullptr && k->parent->color == RED) {
        if (k->parent == k->parent->parent->right) {
            u = k->parent->parent->left;
            if (u->color == RED) {
                u->color = BLACK;
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                k = k->parent->parent;
            } else {
                if (k == k->parent->left) {
                    k = k->parent;
                    rightRotate(k);
                }
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                leftRotate(k->parent->parent);
            }
        } else {
            u = k->parent->parent->right;
            if (u->color == RED) {
                u->color = BLACK;
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                k = k->parent->parent;
            } else {
                if (k == k->parent->right) {
                    k = k->parent;
                    leftRotate(k);
                }
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                rightRotate(k->parent->parent);
            }
        }
        if (k == root) {
            break;
        }
    }
    root->color = BLACK;
}

template<typename T>
void RedBlackTree<T>::remove(T key) {
    Node<T>* z = searchHelper(root, key);
    if (z == nil) return;
    
    Node<T>* y = z;
    Color y_original_color = y->color;
    Node<T>* x;
    
    if (z->left == nil) {
        x = z->right;
        transplant(z, z->right);
    } else if (z->right == nil) {
        x = z->left;
        transplant(z, z->left);
    } else {
        y = minimum(z->right);
        y_original_color = y->color;
        x = y->right;
        if (y->parent == z) {
            x->parent = y;
        } else {
            transplant(y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        transplant(z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    
    delete z;
    if (y_original_color == BLACK) {
        deleteFixup(x);
    }
}

template<typename T>
void RedBlackTree<T>::deleteFixup(Node<T>* x) {
    Node<T>* w;
    while (x != root && x->color == BLACK) {
        if (x == x->parent->left) {
            w = x->parent->right;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                leftRotate(x->parent);
                w = x->parent->right;
            }
            if (w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                if (w->right->color == BLACK) {
                    w->left->color = BLACK;
                    w->color = RED;
                    rightRotate(w);
                    w = x->parent->right;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->right->color = BLACK;
                leftRotate(x->parent);
                x = root;
            }
        } else {
            w = x->parent->left;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                rightRotate(x->parent);
                w = x->parent->left;
            }
            if (w->right->color == BLACK && w->left->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                if (w->left->color == BLACK) {
                    w->right->color = BLACK;
                    w->color = RED;
                    leftRotate(w);
                    w = x->parent->left;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->left->color = BLACK;
                rightRotate(x->parent);
                x = root;
            }
        }
    }
    x->color = BLACK;
}

template<typename T>
void RedBlackTree<T>::transplant(Node<T>* u, Node<T>* v) {
    if (u->parent == nullptr) {
        root = v;
    } else if (u == u->parent->left) {
        u->parent->left = v;
    } else {
        u->parent->right = v;
    }
    v->parent = u->parent;
}

template<typename T>
Node<T>* RedBlackTree<T>::minimum(Node<T>* node) {
    while (node->left != nil) {
        node = node->left;
    }
    return node;
}

template<typename T>
bool RedBlackTree<T>::search(T key) const {
    return searchHelper(root, key) != nil;
}

template<typename T>
Node<T>* RedBlackTree<T>::searchHelper(Node<T>* node, T key) const {
    if (node == nil || key == node->data) {
        return node;
    }
    if (key < node->data) {
        return searchHelper(node->left, key);
    }
    return searchHelper(node->right, key);
}

template<typename T>
void RedBlackTree<T>::print() const {
    inorderHelper(root);
    std::cout << std::endl;
}

template<typename T>
void RedBlackTree<T>::inorderHelper(Node<T>* node) const {
    if (node != nil) {
        inorderHelper(node->left);
        std::cout << node->data << "(" << (node->color == RED ? "R" : "B") << ") ";
        inorderHelper(node->right);
    }
}

template<typename T>
void RedBlackTree<T>::clear() {
    clearHelper(root);
    root = nil;
}

template<typename T>
void RedBlackTree<T>::clearHelper(Node<T>* node) {
    if (node != nil) {
        clearHelper(node->left);
        clearHelper(node->right);
        delete node;
    }
}

#endif 